#include <stdio.h>
#include <stdlib.h>

/* 使用libc自带的free函数 */
void zlibc_free(void *ptr) {
    free(ptr);
}

#include <string.h>
#include <pthread.h>
#include "zmalloc.h"


#ifdef HAVE_MALLOC_SIZE
    #define PREFIX_SIZE (0)
#else
    #define PREFIX_SIZE (sizeof(size_t))
#endif

/* 由于在计算内存时存在竞争, 需要锁机制进行并发控制 */

#ifdef HAVE_ATOMIC
    /* Linux自带的并发 */
    #define update_zmalloc_stat_add(__n) __sync_add_and_fetch(&used_memory, (__n))
    #define update_zmalloc_stat_sub(__n) __sync_sub_and_fetch(&used_memory, (__n))
#else
    /* 基于pthread实现并发 */
    #define update_zmalloc_stat_add(__n) do { \
        pthread_mutex_lock(&used_memory_mutex); \
        used_memory += (__n); \
        pthread_mutex_unlock(&used_memory_mutex); \
    } while (0)

    #define update_zmalloc_stat_sub(__n) do { \
        pthread_mutex_lock(&used_memory_mutex); \
        used_memory -= (__n); \
        pthread_mutex_unlock(&used_memory_mutex); \
    } while (0)

#endif

/* 采用8字节补全的方式进行分配内存 */
#define update_zmalloc_stat_alloc(__n) do { \
    size_t _n = __n; \
    if (_n & (sizeof(long) - 1)) _n += sizeof(long) - (_n & (sizeof(long) - 1)); \
    if (zmalloc_thread_safe) { \
        update_zmalloc_stat_add(_n); \
    } else { \
        used_memory += _n; \
    } \
} while (0)

/* 采用8字节补全的方式释放内存 */
#define update_zmalloc_stat_free(__n) do { \
    size_t _n = __n; \
    if (_n & (sizeof(long) - 1)) _n += sizeof(long) - (_n & (sizeof(long) - 1)); \
    if (zmalloc_thread_safe) { \
        update_zmalloc_stat_sub(_n); \
    } else { \
        used_memory -= _n; \
    } \
} while (0)

static size_t used_memory = 0; /* 使用的内存量, 记录为数值上的使用量 */
static int zmalloc_thread_safe = 0; /* 是否线程安全: 1表示不安全, 0表示安全 */
pthread_mutex_t used_memory_mutex = PTHREAD_MUTEX_INITIALIZER; /* 锁 */

/* oom处理函数 */
static void zmalloc_default_oom(size_t size) {
    fprintf(stderr, "zmalloc: Out of memory trying to allocate %zu bytes\n", size);
    fflush(stderr);
    abort();
}

/* 设置默认的OOM处理函数 */
static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom;


/* 分配PREFIX_SIZE + SIZE内存, 其中PREFIX_SIZE用来保存分配的内存大小 */
void *zmalloc(size_t size) {
    void *ptr = malloc(size + PREFIX_SIZE);
    if (!ptr) zmalloc_oom_handler(size);
#ifdef HAVE_MALLOC_SIZE
    update_zmalloc_stat_alloc(zmalloc_size(ptr));
    return ptr;
#else
    /* 前8个字节设置大小 */
    *((size_t*)ptr) = size;
    update_zmalloc_stat_alloc(size + PREFIX_SIZE);
    /* 返回真正的内容内存 */
    return (char*)ptr + PREFIX_SIZE;
#endif
}

/* 逻辑和malloc一致 */
void *zcalloc(size_t size) {
    void *ptr = calloc(1, size + PREFIX_SIZE);
    if (!ptr) zmalloc_oom_handler(size);
#ifdef HAVE_MALLOC_SIZE
    update_zmalloc_stat_alloc(zmalloc_size(ptr));
    return ptr;
#else
    *((size_t*)ptr) = size;
    update_zmalloc_stat_alloc(size + PREFIX_SIZE);
    return (char*)ptr + PREFIX_SIZE;
#endif
}

/* realloc的逻辑就是重新进行分配, 如果为空指针需要进行对应的操作 */
void *zrealloc(void *ptr, size_t size) {
#ifndef HAVE_MALLOC_SIZE
    void *realptr;
#endif
    size_t oldsize;
    void *newptr;
    if (ptr == NULL) return zmalloc(size);
#ifdef HAVE_MALLOC_SIZE
    oldsize = zmalloc_size(ptr);
    newptr = realloc(ptr, size);
    if (!newptr) zmalloc_oom_handler(size);
    update_zmalloc_stat_free(oldsize);
    update_zmalloc_stat_alloc(zmalloc_size(newptr));
    return newptr;
#else
    realptr = (char*)ptr - PREFIX_SIZE;
    oldsize = *((size_t*)realptr);
    newptr = realloc(realptr, size + PREFIX_SIZE);
    if (!newptr) zmalloc_oom_handler(size);

    *((size_t*)newptr) = size;
    update_zmalloc_stat_free(oldsize + PREFIX_SIZE);
    update_zmalloc_stat_alloc(size + PREFIX_SIZE);
    return (char*)newptr + PREFIX_SIZE;
#endif
}

/* 释放内存, 需要先找到realptr的位置 */
void zfree(void *ptr) {
#ifndef HAVE_MALLOC_SIZE
    void *realptr;
    size_t oldsize;
#endif
    if (ptr == NULL) return;
#ifdef HAVE_MALLOC_SIZE
    update_zmalloc_stat_free(zmalloc_size(ptr));
    free(ptr);
#else
    realptr = (char*)ptr - PREFIX_SIZE;
    oldsize = *((size_t*)ptr);
    update_zmalloc_stat_free(oldsize + PREFIX_SIZE);
    free(realptr); 
#endif
}


/* 记录当前ptr的总共内存占用, 注意字节对齐 */
#ifndef HAVE_MALLOC_SIZE
size_t zmalloc_size(void *ptr) {
    if (ptr == NULL) return 0;
    void *realptr = (char*)ptr - PREFIX_SIZE;
    size_t size = *((size_t*)realptr);
    if (size & (sizeof(long) - 1)) size += sizeof(long) - (size & (sizeof(long) - 1));
    return size + PREFIX_SIZE;
}
#endif

/* 这里相当于只处理了值复制, 但是并没有处理长度计算 */
char *zstrdup(const char *s) {
    size_t l = strlen(s) + 1;
    char *p = zmalloc(l);
    memcpy(p, s, l);
    /* 自行添加的设置字符串占据的内存大小 */
    *((size_t*)(p - PREFIX_SIZE)) = l;
    return p;
}

/* 计算总共使用的内存量, 需要原子操作 */
size_t zmalloc_used_memory(void) {
    size_t um;
    /* 线程不安全 */
    if (zmalloc_thread_safe) {
#ifdef HAVE_ATOMIC
        um = __sync_add_and_fetch(&used_memory, 0);
#else
        pthread_mutex_lock(&used_memory_mutex);
        um = used_memory;
        pthread_mutex_unlock(&used_memory_mutex);
#endif
    } else {
        um = used_memory;
    }
    return um;
}

/* 表示当前线程不安全 */
void zmalloc_enable_thread_safeness(void) {
    zmalloc_thread_safe = 1;
}

/* 设置oom句柄 */
void zmalloc_set_oom_handler(void (*oom_handler)(size_t)) {
    zmalloc_oom_handler = oom_handler;
}

/* 获取碎片比例, 以及后续待做 */
float zmalloc_get_fragmentation_ratio(void) {
    return 1.0f;
}

size_t zmalloc_get_rss(void) {
    return 1;
}
size_t zmalloc_get_private_dirty(void) {
    return 1;
}



